#include <Wire.h>
#include <Rotary.h>
#include <si5351.h>
#include <U8g2lib.h>

// Pin definitions
#define PIN_TUNESTEP A0
#define PIN_BAND     A1
#define PIN_RX_TX    A2
#define PIN_ADC      A3
#define PIN_ROT_1    2
#define PIN_ROT_2    3
#define PIN_RST      8
#define PIN_CS       10
#define PIN_MOSI     11
#define PIN_SCK      13

// Constants
#define IF_FREQ    455
#define BAND_INIT  7
#define XT_CAL_F   33000
#define S_GAIN     303

// Frequency range limits
const uint32_t MIN_FREQ = 10000UL;      // 10 kHz
const uint32_t MAX_FREQ = 225000000UL;  // 225 MHz

// Band names stored in program memory
const char BAND_0[] PROGMEM = " GEN";
const char BAND_1[] PROGMEM = " MW";
const char BAND_2[] PROGMEM = " 160m";
const char BAND_3[] PROGMEM = " 80m";
const char BAND_4[] PROGMEM = " 60m";
const char BAND_5[] PROGMEM = " 49m";
const char BAND_6[] PROGMEM = " 40m";
const char BAND_7[] PROGMEM = " 31m";
const char BAND_8[] PROGMEM = " 25m";
const char BAND_9[] PROGMEM = " 22m";
const char BAND_10[] PROGMEM = " 20m";
const char BAND_11[] PROGMEM = " 19m";
const char BAND_12[] PROGMEM = " 16m";
const char BAND_13[] PROGMEM = " 13m";
const char BAND_14[] PROGMEM = " 11m";
const char BAND_15[] PROGMEM = " 10m";
const char BAND_16[] PROGMEM = " 6m";
const char BAND_17[] PROGMEM = " WFM";
const char BAND_18[] PROGMEM = " AIR";
const char BAND_19[] PROGMEM = " 2m";
const char BAND_20[] PROGMEM = " 1m";

const char* const BAND_NAMES[] PROGMEM = {
  BAND_0, BAND_1, BAND_2, BAND_3, BAND_4, BAND_5, BAND_6, BAND_7, BAND_8, BAND_9,
  BAND_10, BAND_11, BAND_12, BAND_13, BAND_14, BAND_15, BAND_16, BAND_17,
  BAND_18, BAND_19, BAND_20
};

// Frequency presets stored in program memory
const uint32_t FREQ_PRESETS[] PROGMEM = {
  100000UL,    // GEN
  800000UL,    // MW
  1800000UL,   // 160m
  3650000UL,   // 80m
  4985000UL,   // 60m
  6180000UL,   // 49m
  7200000UL,   // 40m
  10000000UL,  // 31m
  11780000UL,  // 25m
  13630000UL,  // 22m
  14100000UL,  // 20m
  15000000UL,  // 19m
  17655000UL,  // 16m
  21525000UL,  // 13m
  27015000UL,  // 11m
  28400000UL,  // 10m
  50000000UL,  // 6m
  100000000UL, // WFM
  130000000UL, // AIR
  144000000UL, // 2m
  220000000UL  // 1m
};

// Frequency steps
const uint32_t FREQ_STEPS[] PROGMEM = {
  1000000UL,  // 1 MHz
  1UL,        // 1 Hz
  10UL,       // 10 Hz
  1000UL,     // 1 kHz
  5000UL,     // 5 kHz
  10000UL     // 10 kHz
};

// Object initialization
U8G2_ST7920_128X64_1_SW_SPI u8g2(U8G2_R0, PIN_SCK, PIN_MOSI, PIN_CS, PIN_RST);
Rotary r = Rotary(PIN_ROT_1, PIN_ROT_2);
Si5351 si5351;

// Global variables
uint32_t freq = 7200000UL;      // Start at 7.2MHz
uint32_t freqold;
uint32_t fstep = 1000;          // Default step 1kHz
int16_t interfreq = IF_FREQ;
int16_t cal = XT_CAL_F;
uint8_t smval;
uint8_t encoder = 1;
uint8_t stp = 4;
uint8_t n = 1;
uint8_t count = BAND_INIT;
uint8_t prevCount = BAND_INIT;
uint8_t x, xo;
bool sts = 0;
bool displayOK = false;

// Function prototypes
bool setSi5351Frequency(Si5351& si5351, uint32_t freq, int16_t interfreq);
void check_inputs();
void update_display_paged();
void initializeSi5351();

// Encoder interrupt service routine
ISR(PCINT2_vect) {
  char result = r.process();
  if (result == DIR_CW) {
    if (encoder == 1) {
      uint32_t new_freq = freq + fstep;
      if (new_freq <= MAX_FREQ) {
        freq = new_freq;
        n = (n >= 42) ? 1 : n + 1;
      }
    }
  }
  else if (result == DIR_CCW) {
    if (encoder == 1) {
      uint32_t new_freq = freq;
      if (freq >= fstep) {
        new_freq = freq - fstep;
        if (new_freq >= MIN_FREQ) {
          freq = new_freq;
          n = (n <= 1) ? 42 : n - 1;
        }
      }
    }
  }
}

void setup() {
  Serial.begin(9600);
  Serial.println(F("VFO Starting..."));
  
  Wire.begin();
  
  if (!u8g2.begin()) {
    Serial.println(F("Display init failed!"));
    while (1) { delay(1000); }
  }
  
  // Display initialization test
  u8g2.setFont(u8g2_font_6x12_tr);
  u8g2.firstPage();
  do {
    u8g2.drawFrame(0, 0, 128, 64);
    u8g2.drawStr(20, 32, "Initializing...");
  } while (u8g2.nextPage());
  delay(1000);
  
  Serial.println(F("Display initialized"));
  displayOK = true;

  // Initialize pins
  pinMode(PIN_ROT_1, INPUT_PULLUP);
  pinMode(PIN_ROT_2, INPUT_PULLUP);
  pinMode(PIN_TUNESTEP, INPUT_PULLUP);
  pinMode(PIN_BAND, INPUT_PULLUP);
  pinMode(PIN_RX_TX, INPUT_PULLUP);

  // Initialize Si5351
  initializeSi5351();
  
  // Setup rotary encoder interrupts
  PCICR |= (1 << PCIE2);
  PCMSK2 |= (1 << PCINT18) | (1 << PCINT19);
  sei();

  // Set initial frequency
  freq = pgm_read_dword(&FREQ_PRESETS[count - 1]);
  
  Serial.println(F("Setup complete"));
}

void initializeSi5351() {
  Serial.println(F("Initializing Si5351..."));
  if (!si5351.init(SI5351_CRYSTAL_LOAD_8PF, 0, 0)) {
    Serial.println(F("Si5351 init failed!"));
  }
  si5351.reset();
  delay(10);
  si5351.set_correction(cal, SI5351_PLL_INPUT_XO);
  si5351.drive_strength(SI5351_CLK0, SI5351_DRIVE_8MA);
  si5351.output_enable(SI5351_CLK0, 1);
}

bool setSi5351Frequency(Si5351& si5351, uint32_t freq, int16_t interfreq) {
  // Check if frequency is within valid range
  if (freq < MIN_FREQ || freq > MAX_FREQ) {
    return false;
  }
  
  uint64_t output_freq = (freq + (interfreq * 1000ULL)) * 100ULL;
  
  // Handle GEN mode specially
  if (count == 1) {
    si5351.reset();
    delay(10);
    si5351.set_correction(cal, SI5351_PLL_INPUT_XO);
    si5351.drive_strength(SI5351_CLK0, SI5351_DRIVE_8MA);
  }
  
  // Set the frequency
  si5351.set_freq(output_freq, SI5351_CLK0);
  si5351.output_enable(SI5351_CLK0, 1);
  
  return true;
}

void loop() {
  if (!displayOK) return;

  // Process frequency changes with error handling
  if (freqold != freq) {
    if (!setSi5351Frequency(si5351, freq, interfreq)) {
      // If frequency setting fails, try to recover
      si5351.reset();
      delay(10);
      si5351.set_correction(cal, SI5351_PLL_INPUT_XO);
      si5351.drive_strength(SI5351_CLK0, SI5351_DRIVE_8MA);
      setSi5351Frequency(si5351, freq, interfreq);
    }
    freqold = freq;
  }

  // Check inputs
  check_inputs();
  
  // Update display
  update_display_paged();
  
  // Read signal meter
  smval = analogRead(PIN_ADC);
  x = constrain(map(smval, 0, S_GAIN, 1, 14), 1, 14);
}

void check_inputs() {
  if (digitalRead(PIN_TUNESTEP) == LOW) {
    stp = (stp % 6) + 1;
    fstep = pgm_read_dword(&FREQ_STEPS[stp - 1]);
    delay(300);
  }

  if (digitalRead(PIN_BAND) == LOW) {
    uint8_t newCount = (count % 21) + 1;
    
    // Reset Si5351 when entering or leaving GEN mode
    if (newCount == 1 || count == 1) {
      si5351.reset();
      delay(10);
      si5351.set_correction(cal, SI5351_PLL_INPUT_XO);
      si5351.drive_strength(SI5351_CLK0, SI5351_DRIVE_8MA);
      si5351.output_enable(SI5351_CLK0, 1);
    }
    
    count = newCount;
    freq = pgm_read_dword(&FREQ_PRESETS[count - 1]);
    prevCount = count;
    delay(300);
  }

  sts = (digitalRead(PIN_RX_TX) == LOW);
  interfreq = (sts || count == 1) ? 0 : IF_FREQ;
}

void update_display_paged() {
  u8g2.firstPage();
  do {
    // Display frequency
    char buffer[16];
    uint32_t m = freq / 1000000UL;
    uint32_t k = (freq % 1000000UL) / 1000UL;
    uint32_t h = (freq % 1000UL);
    
    u8g2.setFont(u8g2_font_10x20_tr);
    
    if (m < 1) {
      sprintf(buffer, "%03lu.%03lu", k, h);
      u8g2.drawStr(41, 17, buffer);
    } else if (m < 100) {
      sprintf(buffer, "%lu.%03lu.%03lu", m, k, h);
      u8g2.drawStr(15, 17, buffer);
    } else {
      sprintf(buffer, "%lu.%03lu.%03lu", m, k, h);
      u8g2.drawStr(15, 17, buffer);
    }
    
    // Draw interface elements
    u8g2.setFont(u8g2_font_6x12_tr);
    u8g2.drawHLine(0, 22, 128);
    u8g2.drawHLine(0, 45, 128);
    u8g2.drawHLine(15, 54, 67);
    u8g2.drawVLine(105, 26, 15);
    u8g2.drawVLine(87, 26, 15);
    u8g2.drawVLine(87, 50, 15);
    
    // Display RX/TX status
    u8g2.drawStr(91, 37, sts ? "TX" : "RX");
    
    // Display IF frequency
    sprintf(buffer, "IF:%d", interfreq);
    u8g2.drawStr(90, 59, buffer);
    
    // Display LO value
    sprintf(buffer, "LO:%d", interfreq);
    u8g2.drawStr(110, 38, buffer);
    
    // Display step
    u8g2.drawStr(54, 32, "STEP");
    switch(stp) {
      case 1: u8g2.drawStr(54, 42, "1MHz"); break;
      case 2: u8g2.drawStr(54, 42, "1Hz"); break;
      case 3: u8g2.drawStr(54, 42, "10Hz"); break;
      case 4: u8g2.drawStr(54, 42, "1kHz"); break;
      case 5: u8g2.drawStr(54, 42, "5kHz"); break;
      case 6: u8g2.drawStr(54, 42, "10kHz"); break;
    }
    
    // Display band name
    u8g2.setFont(u8g2_font_10x20_tr);
    strcpy_P(buffer, (char*)pgm_read_word(&(BAND_NAMES[count - 1])));
    u8g2.drawStr(0, 40, buffer);
    
    // Draw meters
    u8g2.setFont(u8g2_font_6x12_tr);
    byte y = map(n, 1, 42, 1, 14);
    
    u8g2.drawStr(0, 54, "TU");
    u8g2.drawBox(15 + (y-1)*5, 47, 2, 6);
    
    u8g2.drawStr(0, 63, "SM");
    for (byte i = 1; i <= x; i++) {
      u8g2.drawBox(15 + (i-1)*5, 57, 2, 6);
    }
    
  } while (u8g2.nextPage());
}